import pathlib

import matplotlib.pyplot
import torch

import models
import shared.checkpoint
import shared.torchtools


def plot_checkpoint(checkpoint: pathlib.Path) -> matplotlib.pyplot.Figure:
    states = torch.load(checkpoint)

    model = models.PinnOfNavierStokesEquation()
    model.load_state_dict(states["model"])
    model.eval()

    print(f"lambda1: {float(model.lambda1)}")
    print(f"lambda2: {float(model.lambda2)}")

    figure = matplotlib.pyplot.figure(figsize=(8, 5), dpi=110)

    x = torch.linspace(1, 8, 1000)
    y = torch.linspace(-2, 2, 1000)
    x, y = torch.meshgrid(x, y, indexing="xy")
    t = torch.ones_like(x) * 10

    p = model(torch.flatten(t), torch.flatten(x), torch.flatten(y))[:, 2]
    p = shared.torchtools.reshape_as(p, t, detach=True)

    contour = matplotlib.pyplot.contourf(x, y, p,
                                         levels=80, cmap=matplotlib.pyplot.get_cmap("rainbow"))

    matplotlib.pyplot.title("Pressure")
    matplotlib.pyplot.xlabel("$x$")
    matplotlib.pyplot.ylabel("$y$")
    figure.colorbar(contour, fraction=0.08, pad=0.01)

    return figure


def main():
    checkpoint_manager = shared.checkpoint.CheckpointManager(pathlib.Path("./checkpoints"))
    checkpoint = checkpoint_manager.last_existed()
    assert checkpoint is not None

    figure = plot_checkpoint(checkpoint[1])
    figure.show()
    matplotlib.pyplot.pause(0)


if __name__ == "__main__":
    main()
